/**
 * Trie.java 
 * Author: Godwin Babu
 */
package me.godwin.datastructures;

import java.util.ArrayList;
import java.util.List;
import java.util.LinkedList;

/**
 * Trie data structure implementation to store character sequences.
 * 
 * @author Godwin Babu
 * 
 */
public class Trie {
    private static final int MAX_RETURN_LIMIT = 1000;
    // Head is a special node, not part of the data.
    private final Node head = new Node(' ');

    /**
     * Each character is encapsulated inside this class. The a pointer to the
     * child nodes indicate which characters follow this in a word.
     * 
     * @author Godwin Babu
     * 
     */
    static final class Node {
        private final char a;
        private List<Node> children;
        private boolean isWord = false;
        private int count = 0;

        /**
         * Private constructor. Use the factory method for creating a node.
         */
        private Node(final char a) {
            this(a, false);
        }

        /**
         * Private constructor. Use the factory method for creating a node.
         */
        private Node(final char a, final boolean isWord) {
            this.a = a;
            this.isWord = isWord;
        }

        /**
         * Use this factory method for creating a node.
         */
        static Node createNode(final char a, final boolean isWord,
                final Node parent) {
            if (parent == null)
                throw new IllegalArgumentException();
            Node node = parent.addChild(a);
            if (!node.isWord)
                node.isWord = isWord;
            return node;
        }

        Node addChild(final char a) {
            Node child = new Node(a);
            if (children == null) {
                children = new LinkedList<Node>();
                child.count = 1;
                children.add(child);
            } else {
                int ind = children.indexOf(child);
                if (ind < 0) { // not present
                    child.count = 1;
                    children.add(child);
                } else {
                    child = children.get(ind);
                    child.count++;
                }
            }
            return child;
        }

        boolean removeChild(final Node node) {
            if (this.children == null || !this.children.contains(node))
                return false;
            node.count--;
            if (node.count <= 0) {
                assert (node.children == null || node.children.isEmpty()); // invariant
                return this.children.remove(node);
            }
            return true;
        }

        Node getChild(final char a) {
            if (this.children == null)
                return null;
            int ind = this.children.indexOf(new Node(a));
            if (ind >= 0) {
                return this.children.get(ind);
            }
            return null;
        }

        boolean isExists(final char a) {
            return true;
        }

        @Override
        public int hashCode() {
            return (int) a;
        }

        @Override
        public boolean equals(final Object obj) {
            if (obj instanceof Node) {
                return this.a == ((Node) obj).a;
            }
            return false;
        }
    }

    public boolean addWord(final String word) {
        if (word == null)
            return false;
        int len = word.length();
        Node parent = head;
        for (int i = 0; i < len; i++) {
            Node node = Node.createNode(word.charAt(i), i == len - 1, parent);
            parent = node;
        }
        return true;
    }

    public boolean checkWord(final String word) {
        if (word == null)
            return false;
        int len = word.length();
        Node parent = head;
        for (int i = 0; i < len; i++) {
            Node child = parent.getChild(word.charAt(i));
            if (child == null) { // prefix not found
                return false;
            }
            parent = child;
        }
        if (parent != null && parent.isWord) {
            return true;
        }
        return false;
    }

    public boolean removeWord(final String word) {
        if (word == null || word.length() < 1)
            return false;
        int len = word.length();
        Node parent = head;
        LinkedList<Node> queue = new LinkedList<Node>();
        queue.add(parent);
        for (int i = 0; i < len; i++) {
            Node child = parent.getChild(word.charAt(i));
            if (child == null) { // word not found
                return false;
            }
            queue.add(parent);
            parent = child;
        }
        if (parent != null) {
            queue.add(parent); // add the last one
        }

        // the word exists, now remove by polling the queue
        Node child = queue.pollLast();
        while (child != null) {
            parent = queue.pollLast();
            if (parent != null && parent.children != null) { // just in case
                parent.removeChild(child);
            }
            child = parent;
        }
        return true;
    }

    /**
     * Find and return all suffixes for a given prefix
     * 
     * @param prefix
     * @return list of suffix strings, if prefix matches. null, if prefix does
     *         not match.
     * @throws Trie.TooManyMatchesException
     *             when the result size exceed a predefined limit
     */
    public List<String> listSuffix(final String prefix)
            throws Trie.TooManyMatchesException {
        if (prefix == null)
            return null;
        int len = prefix.length();
        Node parent = head;
        for (int i = 0; i < len; i++) {
            Node child = parent.getChild(prefix.charAt(i));
            if (child == null) { // prefix not found
                return null;
            }
            parent = child;
        }

        if (parent == null || parent.children == null
                || parent.children.size() < 1) { // no suffix
            return new ArrayList<String>();
        }

        return listWords(parent);
    }

    /**
     * Return a list of all words in this Trie
     * @return
     * @throws Trie.TooManyMatchesException
     */
    public List<String> listWords() throws Trie.TooManyMatchesException {
        return listWords(head);
    }

    /**
     * Return a list of all words starting from a particular node.
     * @param fromNode
     * @return
     * @throws Trie.TooManyMatchesException
     */
    private List<String> listWords(Node fromNode)
            throws Trie.TooManyMatchesException {
        if (fromNode == null || fromNode.children == null)
            return null;
        LinkedList<Node> queue = new LinkedList<Node>();
        LinkedList<Node> refQueue = 
                new LinkedList<Node>(); // A second queue to backtrack
        StringBuilder buf = new StringBuilder();
        List<String> matches = new ArrayList<String>();
        int i = 0;
        for (Node n : fromNode.children) {
            queue.add(n);
            while (queue.size() > 0) {
                Node node = queue.peek();
                refQueue.push(node);
                if (buf == null) {
                    buf = new StringBuilder();
                }
                buf.append(node.a);
                if (node.isWord) {
                    matches.add(buf.toString());
                    // throw exception if there are too many matches
                    // this will preserve sanity
                    if (++i > MAX_RETURN_LIMIT) {
                        throw new Trie.TooManyMatchesException();
                    }
                }
                if (node.children != null && node.children.size() > 0) {
                    queue.addAll(0, node.children);
                } else {
                    // backtrack
                    queue.remove();
                    refQueue.remove();
                    int numChars = 1;
                    while (!refQueue.isEmpty()) {
                        if (refQueue.peek() == queue.peek()) {
                            refQueue.remove();
                            queue.remove();
                            numChars++;
                        } else {
                            break;
                        }
                    }
                    int bufLen = buf.length();
                    int start = ((bufLen - numChars) >= 0 ? bufLen - numChars
                            : 0);
                    buf.delete(start, buf.length());
                }
            }
        }
        return matches;
    }

    /**
     * Print all words in Trie to System console.
     */
    public void printWords() {
        if (head == null || head.children == null)
            return;
        for (Node n : head.children)
            printWords(n, null);
    }

    /**
     * Print all words in Trie to System console.
     * Depth first traversal using recursion.
     * 
     * @param node
     * @param buf
     */
    private void printWords(final Node node, StringBuilder buf) {
        if (node == null)
            return;
        if (buf == null)
            buf = new StringBuilder();
        buf.append(node.a);
        if (node.isWord) {
            System.out.println(buf.toString());
        }
        if (node.children != null && node.children.size() > 0) {
            for (Node n : node.children) {
                printWords(n, buf);
                buf = buf.deleteCharAt(buf.length() - 1);
            }
        }
        return;
    }

    /**
     * Print all words in Trie to System console. 
     * Depth first traversal using iteration.
     * 
     */
    public void printWords1() {
        if (head == null || head.children == null)
            return;
        LinkedList<Node> queue = new LinkedList<Node>();
        LinkedList<Node> refQueue = 
                new LinkedList<Node>(); // A second queue to backtrack
        StringBuilder buf = new StringBuilder();
        for (Node n : head.children) {
            queue.add(n);
            while (queue.size() > 0) {
                Node node = queue.peek();
                refQueue.push(node);
                if (buf == null) {
                    buf = new StringBuilder();
                }
                buf.append(node.a);
                if (node.isWord) {
                    System.out.println(buf);
                }
                if (node.children != null && node.children.size() > 0) {
                    queue.addAll(0, node.children);
                } else {
                    // backtrack
                    queue.remove();
                    refQueue.remove();
                    int numChars = 1;
                    while (!refQueue.isEmpty()) {
                        if (refQueue.peek() == queue.peek()) {
                            refQueue.remove();
                            queue.remove();
                            numChars++;
                        } else {
                            break;
                        }
                    }
                    int bufLen = buf.length();
                    int start = ((bufLen - numChars) >= 0 ? bufLen - numChars
                            : 0);
                    buf.delete(start, buf.length());
                }
            }
        }
    }

    public static class TooManyMatchesException extends Exception {
        private static final long serialVersionUID = 4803427464041547087L;
    }
}
